# -*- coding: utf-8 -*-
"""
Created on Sun Oct 11 10:09:35 2020
用于库伦耦合的半导体双量子点模型
128初始态对制备保真度

@author: Waikikilick
"""

import numpy as np
from scipy.linalg import expm
from time import *
import multiprocessing as mp
import copy

np.random.seed(1)
I = np.matrix(np.identity(2, dtype=complex))
sx = np.mat([[0, 1], [1, 0]], dtype=complex)
sy = np.mat([[0, -1j], [1j, 0]], dtype=complex)
sz = np.mat([[1, 0], [0, -1]], dtype=complex)
h_1 = 1
h_2 = 1
coupling = 1/2
T = 10*np.pi
dt = np.pi/2
step_max = T/dt
# init_psi = np.mat([[1], [0], [0], [0]], dtype=complex)
# target_psi = np.mat([[1], [0], [0], [1]], dtype=complex)/np.sqrt(2) #最终目标态
init_set_num = 512
target_set_num = init_set_num

action_space =np.array([[1,1],
                        [1,2],
                        [1,3],
                        [1,4],
                        [1,5],
                        [2,1],
                        [2,2],
                        [2,3],
                        [2,4],
                        [2,5],
                        [3,1],
                        [3,2],
                        [3,3],
                        [3,4],
                        [3,5],
                        [4,1],
                        [4,2],
                        [4,3],
                        [4,4],
                        [4,5],
                        [5,1],
                        [5,2],
                        [5,3],
                        [5,4],
                        [5,5]
                        ])
n_actions = len(action_space)

alpha_num = 4
theta = [np.pi/8,np.pi/4,3*np.pi/8]
theta_1 = theta
theta_2 = theta
theta_3 = theta

alpha = np.linspace(0,np.pi*2,alpha_num,endpoint=False)
alpha_1 = alpha
alpha_2 = alpha
alpha_3 = alpha
alpha_4 = alpha

a_list_complex = np.matrix([[0,0,0,0]],dtype=complex) #第一行用来占位，否则无法和其他行并在一起，在最后要注意去掉这一行
for ii in range(3): #theta_1
    for jj in range(3): #theta_2
        for kk in range(3): #theta_3
            for mm in range(alpha_num): #alpha_1
                for nn in range(alpha_num): #alpha_2
                    for oo in range(alpha_num): #alpha_3
                        for pp in range(alpha_num): #alpha_4
                            
                            a_1_mo = np.cos(theta_1[ii])
                            a_2_mo = np.sin(theta_1[ii])*np.cos(theta_2[jj])
                            a_3_mo = np.sin(theta_1[ii])*np.sin(theta_2[jj])*np.cos(theta_3[kk])
                            a_4_mo = np.sin(theta_1[ii])*np.sin(theta_2[jj])*np.sin(theta_3[kk])
                            
                            a_1_real = a_1_mo*np.cos(alpha_1[mm])
                            a_1_imag = a_1_mo*np.sin(alpha_1[mm])
                            a_2_real = a_2_mo*np.cos(alpha_2[nn])
                            a_2_imag = a_2_mo*np.sin(alpha_2[nn])
                            a_3_real = a_3_mo*np.cos(alpha_3[oo])
                            a_3_imag = a_3_mo*np.sin(alpha_3[oo])
                            a_4_real = a_4_mo*np.cos(alpha_4[pp])
                            a_4_imag = a_4_mo*np.sin(alpha_4[pp])
                            
                            a_1_complex = a_1_real + a_1_imag*1j
                            a_2_complex = a_2_real + a_2_imag*1j
                            a_3_complex = a_3_real + a_3_imag*1j
                            a_4_complex = a_4_real + a_4_imag*1j
                            
                            a_complex = np.matrix([[ a_1_complex, a_2_complex, a_3_complex, a_4_complex]])
                            a_list_complex = np.row_stack((a_list_complex,a_complex))
                            
psi_set = np.array(np.delete(a_list_complex,0,axis=0)) # 删除矩阵的第一行
np.random.shuffle(psi_set) #打乱顺序
init_set = psi_set[:init_set_num]

def target_set():
    target_set = psi_set[init_set_num : init_set_num + target_set_num]
    return target_set

#动作直接选最优的
def step0(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        J_1, J_2=  action_space[action,0], action_space[action,1]  # control field strength
        J_12 = J_1 * J_2 /2
    
        H =  (J_1*np.kron(sz, I) + J_2*np.kron(I, sz) + \
                        J_12/2*np.kron((sz - I),(sz - I)) + \
           h_1*np.kron(sx,I) + h_2*np.kron(I,sx))*coupling
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_


#动作选最优的，或者最差的
def step1(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        J_1, J_2=  action_space[action,0], action_space[action,1]  # control field strength
        J_12 = J_1 * J_2 /2
    
        H =  (J_1*np.kron(sz, I) + J_2*np.kron(I, sz) + \
                        J_12/2*np.kron((sz - I),(sz - I)) + \
           h_1*np.kron(sx,I) + h_2*np.kron(I,sx))*coupling

        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
    
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        
        best_action = fid_list.index(min(fid_list))
        best_fid = min(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_

#动作选最优的，或者次优的
def step2(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        J_1, J_2=  action_space[action,0], action_space[action,1]  # control field strength
        J_12 = J_1 * J_2 /2
    
        H =  (J_1*np.kron(sz, I) + J_2*np.kron(I, sz) + \
                        J_12/2*np.kron((sz - I),(sz - I)) + \
           h_1*np.kron(sx,I) + h_2*np.kron(I,sx))*coupling
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
        
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        psi_list_ = copy.deepcopy(psi_list)
        fid_list_ = copy.deepcopy(fid_list)
        
        del psi_list_[fid_list_.index(max(fid_list_))]
        del fid_list_[fid_list_.index(max(fid_list_))]
        
        best_action = fid_list.index(max(fid_list_))
        
        best_fid = max(fid_list_)
        
    psi_ = psi_list[best_action]
    
    return best_action, best_fid, psi_
#---------------------------------------------------------------------------------
#将数值列表从小到大进行排序（冒泡法）
def sort_fid_list(lt):
    n= len(lt)
    for x in range(n-1):
       for y in range(n-1-x):
          if lt[y]>lt[y+1]:
             lt[y],lt[y+1]=lt[y+1],lt[y]
    return lt
#--------------------------------------------------------------------------------




time1 = time()
def job(target_psi):
    fids_list = []
    target_psi = np.mat(target_psi).T
    for k in range(init_set_num):
        psi1 = np.mat(init_set[k]).T
        psi = psi1
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real
        
        fid_max = F
        fid_max1 = F
        fid_max2 = F
        fid_max0 = F
        
        step_n = 0
        while True:
            action, F, psi_ = step1(psi,target_psi,F)
            
            fid_max1 = max(F,fid_max1)
            psi = psi_
            step_n += 1
            if fid_max1>0.999 or step_n>step_max:
                break
            
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        while True:
            action, F, psi_ = step2(psi,target_psi,F)
            fid_max2 = max(F,fid_max2)
            psi = psi_
            step_n += 1
            if fid_max2>0.999 or step_n>step_max:
                break 
            
    
       
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        while True:
            action, F, psi_ = step0(psi,target_psi,F)
            fid_max0 = max(F,fid_max0)
            psi = psi_
            step_n += 1
            if fid_max0>0.999 or step_n>step_max:
                break 
            
        fid_max = max(fid_max1,fid_max2,fid_max0)  
        fids_list.append(fid_max)
        
    return  np.mean(fids_list)


def multicore():
    pool = mp.Pool()
    F_list = pool.map(job, target_set)
    return F_list
    
if __name__ == '__main__':

    target_set = target_set()
    # target_set = [target_set[48]]
    time1 = time()
    F_list = multicore()
    print(F_list)
    print(np.mean(F_list))
    time2 = time()
    print('time_cost is: ',time2-time1)

[0.97303214885914, 0.8581171564814102, 0.8486689597028498, 0.9520563280326086, 0.8877394014416216, 0.9258675470389715, 0.935394027379221, 0.9423640095554561, 0.9372116412589162, 0.9038594027405402, 0.8648418040080401, 0.8871346325571161, 0.8884062257265531, 0.8808062168253923, 0.9348714962189144, 0.9393246420569237, 0.9147016836557317, 0.8617920783917083, 0.918301937562628, 0.9077757362797709, 0.9606976272620233, 0.920363531424637, 0.8902086413246988, 0.9296976003079437, 0.9067477099028333, 0.8687336883264509, 0.9051528561754703, 0.8648335583047342, 0.9011944971150059, 0.9434366353639448, 0.9178917690935606, 0.8869054329868931, 0.9181287647093835, 0.9365514294997288, 0.9298386251292552, 0.9025915522147526, 0.8897998133511283, 0.8726731254128379, 0.8324651235631612, 0.8460304387871007, 0.9197066986078306, 0.8870186045277038, 0.9609308112098116, 0.9671475525542734, 0.9485847495862848, 0.9141581350317092, 0.9218583141093933, 0.949295796937659, 0.7795960656171885, 0.9453207240323767, 0.8543215878557691, 0.871271335479312, 0.9507106925206585, 0.9692395886525411, 0.9161936236899122, 0.8975747534815162, 0.9272131222648254, 0.9423433200779531, 0.9691485903443822, 0.9458329871208663, 0.870549279873546, 0.9126094900533077, 0.9248400923409923, 0.8973129342204718, 0.8471235296789296, 0.9525343737308714, 0.9581937571032157, 0.8366417870551195, 0.9105626696711545, 0.9203166349309424, 0.9006169993524356, 0.9178182073636971, 0.935598037565279, 0.9218583141093935, 0.9150174188629036, 0.8629399564708653, 0.916922013751316, 0.9268151689144388, 0.8757395360442595, 0.864541901255162, 0.9443682575441347, 0.8759811196902584, 0.9066185891442904, 0.8734114106095889, 0.9153663248388002, 0.9273014930831287, 0.9234162743221412, 0.9411341086882298, 0.8438484994655009, 0.9387549637277834, 0.8782870570041889, 0.9338184980418144, 0.8875341188151741, 0.9179489927787019, 0.8116920445689924, 0.9196270697782278, 0.8875341188151741, 0.9707199104997949, 0.8757140695944634, 0.9640362523852085, 0.9431243170276813, 0.9200274203779116, 0.8590771580900727, 0.8998384376690842, 0.9063873420798729, 0.9414469965088474, 0.9433020279284303, 0.9006169993524357, 0.9164713973765193, 0.8607441019258819, 0.8429779176624816, 0.9068893789401299, 0.9166980257558838, 0.8516389343324456, 0.8862552188713719, 0.9637532111725055, 0.830566036724625, 0.9066185891442904, 0.8907584155564834, 0.8624563106461919, 0.9079676033862465, 0.9293086848752063, 0.9514567820324267, 0.9846236362982921, 0.8941244887665809, 0.9029082025550819, 0.9034339771373519, 0.9611764898773464, 0.8961493192767195, 0.9208726708764571, 0.9092116005137358, 0.934528971247751, 0.8859585763953374, 0.9195681951848502, 0.877936598119752, 0.8983653448447267, 0.9326277119261522, 0.9176267783936604, 0.9342299087063866, 0.8778204852931377, 0.8356548974754511, 0.8946172579289428, 0.8551906090552719, 0.9162153081902319, 0.9055871592303966, 0.9547499381121658, 0.8493990791134666, 0.9221610839555248, 0.9106269011194188, 0.9525108190133642, 0.9525343737308714, 0.9113588790415881, 0.8855429330511523, 0.9502491016418244, 0.9167006336952349, 0.8778269095719948, 0.9611764898773463, 0.8905564743666535, 0.930542375267083, 0.9063677459128479, 0.864541901255162, 0.8831567349336626, 0.930900312615852, 0.9230527658161385, 0.8350571573114277, 0.9297054847365518, 0.91556099496435, 0.923531656342871, 0.9249542607288199, 0.9458329871208664, 0.9239228355975337, 0.9460750682308605, 0.8338297040355886, 0.8442860734892458, 0.9176933364186775, 0.919537434469436, 0.968203559285431, 0.9119477044818509, 0.9810262469939237, 0.8789230587328665, 0.9091649210801533, 0.8705958235126735, 0.8888937772858084, 0.8921872187971532, 0.9117282675345666, 0.9416272758851639, 0.877936598119752, 0.9381146499761615, 0.931428045900373, 0.9276287809905405, 0.9006169993524356, 0.8356548974754511, 0.9009914453854369, 0.9547639706470058, 0.8831690657857724, 0.8386048820904024, 0.9143274992580641, 0.9116335577166055, 0.9529275501017608, 0.8905564743666535, 0.9334898304908649, 0.8555985342083421, 0.9335636587009302, 0.9198697689452959, 0.9201909589907952, 0.8624563106461919, 0.9665115856704585, 0.9298738808630753, 0.8973832748619048, 0.8838692125729726, 0.8725002403172917, 0.8878097169326171, 0.9129059230504956, 0.85680465781009, 0.9240919121292219, 0.8865954074494331, 0.979667022096078, 0.9219263062007749, 0.9277052904962614, 0.8961229081460595, 0.9368746215564592, 0.8778204852931377, 0.9285249217594964, 0.9184813962213961, 0.9439145122946468, 0.9330785332054139, 0.9355291206569452, 0.9480136375492729, 0.907983174339246, 0.8779334219355825, 0.93288840622562, 0.9219263062007749, 0.9145564921797715, 0.909639772121882, 0.8559211632128696, 0.8970099343469979, 0.8930597416137, 0.8700160523639129, 0.9821156452103306, 0.8848750064934228, 0.9034339771373521, 0.9218898369825841, 0.9116340575293609, 0.9330108497514616, 0.8878185231802832, 0.8974905238483282, 0.8875119252074657, 0.9372116412589162, 0.9411341086882298, 0.9532204635150603, 0.959286581759278, 0.9630572500264698, 0.9234350782410468, 0.9272988674764098, 0.9099334046358623, 0.9657752820314993, 0.9642232945828406, 0.8533610065009953, 0.9256445812170978, 0.8734618874133416, 0.8993197775365829, 0.8882062432103512, 0.9380933806245165, 0.9045585619842516, 0.9147151612669207, 0.8533610065009953, 0.8872726995660094, 0.9021306636763269, 0.9938531051473154, 0.9265421637808662, 0.9263310199125148, 0.9019546357484312, 0.9228875435973316, 0.9760226013345619, 0.972108458404644, 0.9315384027422324, 0.9373796276242794, 0.9049875964177025, 0.9075535281035068, 0.9281344899595085, 0.9119294621744618, 0.912041739340659, 0.8520733600155286, 0.9147492060185882, 0.966459853183848, 0.8335429936530233, 0.9243905878728116, 0.9268196210245915, 0.9081103798963737, 0.8700160523639129, 0.9098137984824926, 0.9114521094578454, 0.8929539007707707, 0.9282352473790504, 0.9049875964177024, 0.8471235296789296, 0.9064691758041046, 0.934528971247751, 0.9368794981449415, 0.9117699711213785, 0.9434868694919517, 0.9441299788502734, 0.9076574089594341, 0.9760858708549175, 0.8942925412618583, 0.9618577787444245, 0.9099116775388353, 0.9371316100119043, 0.8776560921208367, 0.9123060277691851, 0.9300377429445152, 0.915129916467197, 0.8680745287324333, 0.9200557927372136, 0.9955474022972512, 0.9134092482960159, 0.9251217184436006, 0.8670228091864259, 0.907654419587929, 0.8713276505271352, 0.9123257771857061, 0.93230801083628, 0.9258489198926635, 0.9197959904177649, 0.8624563106461919, 0.9137799637295687, 0.8942852771755944, 0.9196430904980797, 0.8915865144525157, 0.9669359461541798, 0.9387480274084707, 0.9414469965088474, 0.8937016019004513, 0.9392825761784616, 0.8878436173805924, 0.9140172311125063, 0.9069601486023062, 0.8603165705590455, 0.9116254975510993, 0.9460750682308604, 0.8893991072327577, 0.897220985940427, 0.9253316395374767, 0.9199673878299738, 0.8583694248774696, 0.9013880992546253, 0.9450472148530309, 0.9286445269549495, 0.9189435752504729, 0.8870186045277038, 0.922462756580835, 0.9027409493540772, 0.9163511041025613, 0.8964092992631008, 0.9064691758041046, 0.9401650456258677, 0.8116920445689924, 0.9187401543965819, 0.8852074482122269, 0.9439378235644492, 0.9189626316787568, 0.9263310199125149, 0.8636348245650646, 0.896509661547736, 0.9173161221282563, 0.8872726995660094, 0.9028087499507196, 0.8324651235631612, 0.9147016836557317, 0.9455411479203231, 0.9266167902310247, 0.9082446673924629, 0.9045651408291471, 0.8591657247140281, 0.8660861761244156, 0.9442500887275712, 0.876617261056288, 0.9456918212918048, 0.9047899028530222, 0.9341695068158176, 0.9251217184436006, 0.9231226536081207, 0.873206216960319, 0.9223014238817391, 0.97510041410335, 0.94605881744894, 0.9704551052623553, 0.9085455438958026, 0.8549802286198978, 0.9229833773302747, 0.8607845196239978, 0.918501899537995, 0.9037451419879761, 0.9141766286857207, 0.9121859387539841, 0.9091385324318373, 0.8728970786731893, 0.8957162817153626, 0.8927065097177469, 0.8999206842461303, 0.892558383089536, 0.9162118735389317, 0.9342454448248496, 0.9328790808242384, 0.9514541866749546, 0.917598859640415, 0.9097315741597352, 0.9043860566263762, 0.9018137839307792, 0.9085870958090926, 0.9352829624386485, 0.939816028291326, 0.9305015166525886, 0.8962433230561755, 0.9037305467992908, 0.9044596628110688, 0.953539850681969, 0.9122128511953025, 0.8305660367246249, 0.8822680040465438, 0.8493441360296766, 0.9555125034530805, 0.935598037565279, 0.8968228573945315, 0.9328790808242385, 0.862647708549864, 0.9477767794140797, 0.9063244076215717, 0.9570697876391911, 0.9573979589965957, 0.9678095629255985, 0.9326679892416613, 0.8973536011815111, 0.9099334046358623, 0.9154713339143379, 0.9200250939267184, 0.9148633954811176, 0.9027409493540772, 0.9021306636763269, 0.9050105371689794, 0.8893991072327578, 0.9392825761784616, 0.8494607019866044, 0.9183156891803359, 0.9095721889273825, 0.9414125104374671, 0.7977037847241835, 0.9618577787444245, 0.9177283437183046, 0.9229453138709244, 0.9178182073636971, 0.9502491016418244, 0.9848385134568988, 0.9459087525849865, 0.9365514294997288, 0.9253468701196723, 0.9413277429222482, 0.8993197775365829, 0.9023988935506226, 0.9421567555071988, 0.9466787310214516, 0.9059729061157288, 0.9172755364582723, 0.8640273095723437, 0.876617261056288, 0.9339605692957837, 0.8738698112027636, 0.9203582325242136, 0.9072451619890368, 0.8831444451878163, 0.8913511426580347, 0.939689879413032, 0.8708454934349518, 0.9012142123350774, 0.9342299087063866, 0.9330785332054139, 0.9164903403492091, 0.9327504528866535, 0.9504336275264068, 0.8935793152547066, 0.9170066584407683, 0.9181555840498505, 0.9015252451994811, 0.9320561602615984, 0.8223022582355493, 0.825626637659125, 0.9092607588190328, 0.9147470118514744, 0.8784222976110357, 0.9786778676273671, 0.9106969773054179, 0.8931581114754825, 0.9099334046358623, 0.8915865144525157, 0.8916444655669888, 0.8893685752072367, 0.8834577165350694, 0.9668822002396451, 0.8983653448447266, 0.9055871592303966, 0.9086971142541826, 0.9077015183272221, 0.904630667458967, 0.9612247187853982, 0.9192971106553847, 0.9179150849886601, 0.8341104536720729, 0.9343871985082229, 0.9577699378340871, 0.9411130831410013, 0.9338626777859882, 0.9604405760138534]
0.910852024615862
time_cost is:  20102.19739151001


0.00-0.49:  0
0.50-0.54:  0 0 0 0 0
0.55-0.59:  0 0 0 0 0
0.60-0.64:  0 0 0 0 0
0.65-0.69:  0 0 0 0 0
0.70-0.74:  0 0 0 0 0
0.75-0.79:  0 0 0 1 0
0.80-0.84:  1 2 1 8 8
0.85-0.89:  13 23 18 23 45
0.90-0.94:  51 72 73 57 43
0.95-1.00:  30 18 14 9 1 1
